import numpy as np
from PIL import Image
import torch


class ImagesToSpritesheetNode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "padding": ("INT", {"default": 300, "min": 0, "max": 1000, "step": 10}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "convert_images_to_spritesheet"
    CATEGORY = "image"

    def remove_background(self, image, tolerance=30):
        """
        自动扣除图片背景
        """
        # 转换为RGBA模式以支持透明度
        image = image.convert("RGBA")
        
        # 获取图片数据
        data = np.array(image)
        
        # 获取左上角像素的颜色作为背景色
        bg_color = data[0, 0]
        
        # 创建掩码，标记与背景色相似的像素（在容差范围内）
        # 计算颜色差异
        diff = np.abs(data[:, :, :3].astype(np.int16) - bg_color[:3].astype(np.int16))
        mask = np.all(diff <= tolerance, axis=2)
        
        # 将背景像素设为透明
        data[mask] = [0, 0, 0, 0]
        
        # 创建新图片
        no_bg_image = Image.fromarray(data, 'RGBA')
        
        return no_bg_image

    def convert_images_to_spritesheet(self, images, padding):
        """
        将图像序列转换为精灵表
        """
        # 将tensor或numpy数组转换为PIL图像
        pil_images = []
        
        # 确保我们处理的是numpy数组
        if torch.is_tensor(images):
            # 如果是PyTorch tensor，转换为numpy
            images_np = images.cpu().numpy()
        else:
            # 如果已经是numpy数组
            images_np = images
        
        for i in range(images_np.shape[0]):
            # ComfyUI中的图像格式是 (B, H, W, C) 范围是 0-1
            # 转换为 0-255 的 numpy 数组
            image_array = (images_np[i] * 255).astype(np.uint8)
            pil_image = Image.fromarray(image_array, 'RGB')
            pil_images.append(pil_image)
        
        # 去除帧的背景
        processed_images = []
        for img in pil_images:
            img_no_bg = self.remove_background(img)
            processed_images.append(img_no_bg)
        
        # 计算spritesheet的尺寸
        # 使用网格布局，尽可能接近正方形
        num_images = len(processed_images)
        grid_size = int(np.ceil(np.sqrt(num_images)))
        
        # 计算最大宽度和高度
        max_width = max(img.width for img in processed_images)
        max_height = max(img.height for img in processed_images)
        
        # 添加padding到尺寸计算中
        cell_width = max_width + padding * 2
        cell_height = max_height + padding * 2
        
        # 创建spritesheet
        sheet_width = grid_size * cell_width
        sheet_height = grid_size * cell_height
        spritesheet = Image.new('RGBA', (sheet_width, sheet_height), (0, 0, 0, 0))
        
        # 将图片放置到spritesheet上
        for i, img in enumerate(processed_images):
            row = i // grid_size
            col = i % grid_size
            
            # 计算位置（在单元格内居中放置，包含padding）
            x = col * cell_width + padding + (max_width - img.width) // 2
            y = row * cell_height + padding + (max_height - img.height) // 2
            
            # 粘贴图片
            spritesheet.paste(img, (x, y), img)
        
        # 转换为numpy数组并规范化到0-1范围
        spritesheet_np = np.array(spritesheet).astype(np.float32) / 255.0
        
        # 确保维度顺序正确 (H, W, C)
        # ComfyUI期望的IMAGE格式是 (B, H, W, C)，其中B是批次维度
        # 我们需要添加批次维度，所以最终形状是 (1, H, W, C)
        if len(spritesheet_np.shape) == 3:
            # 添加批次维度
            spritesheet_tensor = np.expand_dims(spritesheet_np, axis=0)
        else:
            # 如果已经是正确的形状
            spritesheet_tensor = spritesheet_np
        
        # 确保返回的是PyTorch tensor而不是numpy array
        # 因为ComfyUI通常期望PyTorch tensors
        if not torch.is_tensor(spritesheet_tensor):
            spritesheet_tensor = torch.from_numpy(spritesheet_tensor)
        
        return (spritesheet_tensor,)

# Node class mappings
NODE_CLASS_MAPPINGS = {
    "ImagesToSpritesheetNode": ImagesToSpritesheetNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ImagesToSpritesheetNode": "Images to Spritesheet"
}